{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os\n",
    "import shutil\n",
    "import time\n",
    "from pathlib import Path\n",
    "\n",
    "import cv2\n",
    "import torch\n",
    "import torch.backends.cudnn as cudnn\n",
    "from numpy import random\n",
    "\n",
    "from models.experimental import attempt_load\n",
    "from utils.datasets import LoadStreams, LoadImages\n",
    "from utils.general import (\n",
    "    check_img_size, non_max_suppression, apply_classifier, scale_coords,\n",
    "    xyxy2xywh, plot_one_box, strip_optimizer, set_logging)\n",
    "from utils.torch_utils import select_device, load_classifier, time_synchronized"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 初始化"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using CUDA device0 _CudaDeviceProperties(name='GeForce GTX 1650', total_memory=4096MB)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "set_logging()\n",
    "device=''\n",
    "device = select_device(device)\n",
    "out='./reference/output'\n",
    "\n",
    "\n",
    "if os.path.exists(out):  # output dir\n",
    "    shutil.rmtree(out)  # delete dir\n",
    "os.makedirs(out)  # make new dir\n",
    "half = device.type != 'cpu'  # half precision only supported on CUDA"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 加载模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fusing layers... \n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Model Summary: 140 layers, 7.45958e+06 parameters, 0 gradients\n"
     ]
    }
   ],
   "source": [
    "weights='yolov5s.pt'\n",
    "imgsz=640\n",
    "\n",
    "\n",
    "model = attempt_load(weights, map_location=device)  # load FP32 model\n",
    "imgsz = check_img_size(imgsz, s=model.stride.max())  # check img_size\n",
    "if half:\n",
    "    model.half()  # to FP16"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Set Dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "source='huge.jpg'\n",
    "\n",
    "vid_path, vid_writer = None, None\n",
    "\n",
    "save_img = True\n",
    "dataset = LoadImages(source, img_size=imgsz)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get names and colors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "names = model.module.names if hasattr(model, 'module') else model.names\n",
    "colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Run inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "image 1/1 D:\\projects\\DL\\1\\minyolo5\\huge.jpg: "
     ]
    }
   ],
   "source": [
    "augment='store_true'\n",
    "conf_thres=0.5\n",
    "iou_thres=0.45\n",
    "classes=0\n",
    "agnostic_nms='store_true'\n",
    "classify=0\n",
    "\n",
    "t0 = time.time()\n",
    "img = torch.zeros((1, 3, imgsz, imgsz), device=device)  # init img\n",
    "_ = model(img.half() if half else img) if device.type != 'cpu' else None  # run once\n",
    "for path, img, im0s, vid_cap in dataset:\n",
    "    img = torch.from_numpy(img).to(device)\n",
    "    img = img.half() if half else img.float()  # uint8 to fp16/32\n",
    "    img /= 255.0  # 0 - 255 to 0.0 - 1.0\n",
    "    if img.ndimension() == 3:\n",
    "        img = img.unsqueeze(0)\n",
    "        # Inference\n",
    "        t1 = time_synchronized()\n",
    "        pred = model(img, augment=augment)[0]\n",
    "        \n",
    "        # Apply NMS\n",
    "        pred = non_max_suppression(pred, conf_thres, iou_thres, classes=classes, agnostic=agnostic_nms)\n",
    "        t2 = time_synchronized()\n",
    "\n",
    "        # Apply Classifier\n",
    "        if classify:\n",
    "            pred = apply_classifier(pred, modelc, img, im0s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Process detections"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "xywh [0.8333333134651184, 0.5, 0.3333333432674408, 1.0]\n",
      "xyxy [tensor(640., device='cuda:0'), tensor(0., device='cuda:0'), tensor(960., device='cuda:0'), tensor(540., device='cuda:0')]\n"
     ]
    }
   ],
   "source": [
    "save_txt=1\n",
    "save_img=1\n",
    "save_conf='store_true'\n",
    "view_img=1\n",
    "\n",
    "\n",
    "for i, det in enumerate(pred):  # detections per image\n",
    "    p, s, im0 = path, '', im0s\n",
    "    save_path = str(Path(out) / Path(p).name)\n",
    "    txt_path = str(Path(out) / Path(p).stem) + ('_%g' % dataset.frame if dataset.mode == 'video' else '')\n",
    "    s += '%gx%g ' % img.shape[2:]  # print string\n",
    "    gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh\n",
    "    if det is not None and len(det):\n",
    "        # Rescale boxes from img_size to im0 size\n",
    "        det[:, :4] = scale_coords(img.shape[2:], det[:, :4], im0.shape).round()\n",
    "        # Print results\n",
    "        for c in det[:, -1].unique():\n",
    "            n = (det[:, -1] == c).sum()  # detections per class\n",
    "            s += '%g %ss, ' % (n, names[int(c)])  # add to string\n",
    "            # Write results\n",
    "            for *xyxy, conf, cls in reversed(det):\n",
    "                if save_txt:  # Write to file\n",
    "                    xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist()  # normalized xywh\n",
    "                    print('xywh',xywh)\n",
    "                    line = (cls, conf, *xywh) if save_conf else (cls, *xywh)  # label format\n",
    "                    with open(txt_path + '.txt', 'a') as f:\n",
    "                        f.write(('%g ' * len(line) + '\\n') % line)\n",
    "                if save_img or view_img:  # Add bbox to image\n",
    "                    label = '%s %.2f' % (names[int(cls)], conf)\n",
    "                    plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3)\n",
    "                    print('xyxy',xyxy)\n",
    "    # Stream results\n",
    "    if view_img:\n",
    "        cv2.imshow(p, im0)\n",
    "        cv2.waitKey(0)\n",
    "        if cv2.waitKey(1) == ord('q'):  # q to quit\n",
    "            raise StopIteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
